"""
流式数据格式配置文件
可以通过修改这个文件来自定义不同的输出格式
"""

from typing import Dict, Any, Callable
from datetime import datetime
import uuid

class StreamFormatConfig:
    """流式数据格式配置类"""
    
    @staticmethod
    def default_format(chunk: Dict[str, Any], message_order: int, **kwargs) -> Dict[str, Any]:
        """默认格式 - 与原始StreamResponse相同"""
        return {
            "type": "message" if chunk["type"] != "error" else "error",
            "content": f"{chunk.get('step', '')}: {chunk['content']}",
            "message_order": message_order,
            "token_count": None,
            "model_name": None,
            "processing_time_ms": None
        }
    
    @staticmethod
    def detailed_format(chunk: Dict[str, Any], message_order: int, **kwargs) -> Dict[str, Any]:
        """详细格式 - 包含更多元数据"""
        return {
            "type": "message",
            "step": chunk.get("step"),
            "content": chunk["content"],
            "chunk_type": chunk["type"],
            "progress": kwargs.get("progress", 0),
            "timestamp": datetime.now().isoformat(),
            "message_order": message_order,
            "metadata": {
                "step": chunk.get("step"),
                "original_type": chunk["type"]
            }
        }
    
    @staticmethod
    def frontend_format(chunk: Dict[str, Any], message_order: int, **kwargs) -> Dict[str, Any]:
        """前端友好格式"""
        # 映射内容类型
        content_type_mapping = {
            "reasoning": "text",
            "element": "text", 
            "ggb_commands": "commands",
            "html_code": "html",
            "optimized_html": "html"
        }
        
        return {
            "id": kwargs.get("message_id", str(uuid.uuid4())),
            "type": content_type_mapping.get(chunk["type"], "text"),
            "payload": chunk["content"],
            "step_name": chunk.get("step"),
            "step_progress": kwargs.get("step_progress", 0),
            "is_final": chunk["type"] == "optimized_html",
            "created_at": datetime.now().isoformat()
        }
    
    @staticmethod
    def simple_format(chunk: Dict[str, Any], message_order: int, **kwargs) -> Dict[str, Any]:
        """简化格式"""
        return {
            "status": "processing",
            "message": chunk["content"],
            "data": {
                "step": chunk.get("step"),
                "type": chunk["type"]
            }
        }
    
    @staticmethod
    def custom_format(chunk: Dict[str, Any], message_order: int, **kwargs) -> Dict[str, Any]:
        """自定义格式 - 你可以在这里定义自己的格式"""
        return {
            "event": "data_chunk",
            "stage": chunk.get("step", "unknown"),
            "content": chunk["content"],
            "meta": {
                "type": chunk["type"],
                "order": message_order,
                "timestamp": datetime.now().isoformat(),
                "conversation_id": kwargs.get("conversation_id")
            }
        }

# 格式注册表
STREAM_FORMATS: Dict[str, Callable] = {
    "default": StreamFormatConfig.default_format,
    "detailed": StreamFormatConfig.detailed_format,
    "frontend": StreamFormatConfig.frontend_format,
    "simple": StreamFormatConfig.simple_format,
    "custom": StreamFormatConfig.custom_format,
}

# 步骤映射配置
STEP_MAPPING = {
    "element_extraction": {
        "name": "元素提取",
        "description": "分析数学问题并提取几何元素",
        "progress": 0.25,
        "order": 1
    },
    "ggb_generation": {
        "name": "GGB命令生成", 
        "description": "根据提取的元素生成GeoGebra命令",
        "progress": 0.5,
        "order": 2
    },
    "html_generation": {
        "name": "HTML生成",
        "description": "将GeoGebra命令转换为HTML代码",
        "progress": 0.75,
        "order": 3
    },
    "html_optimization": {
        "name": "HTML优化",
        "description": "优化生成的HTML代码",
        "progress": 1.0,
        "order": 4
    }
}

def get_format_function(format_name: str) -> Callable:
    """获取指定格式的处理函数"""
    return STREAM_FORMATS.get(format_name, STREAM_FORMATS["default"])

def format_chunk(chunk: Dict[str, Any], format_name: str = "default", **kwargs) -> Dict[str, Any]:
    """格式化数据块"""
    format_func = get_format_function(format_name)
    return format_func(chunk, **kwargs) 